// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include "signatures.h" // @manual
#include <pollux/exec/aggregate.h>
#include <pollux/functions/function_registry.h>
#include <pollux/functions/prestosql/aggregates/register_aggregate_functions.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/functions/sparksql/aggregates/register.h>
#include <pollux/functions/sparksql/registration/register.h>

namespace kumo::pollux::py {

namespace py = pybind11;

void registerPrestoFunctions(const std::string& prefix) {
  kumo::pollux::functions::prestosql::registerAllScalarFunctions(prefix);
}

void registerSparkFunctions(const std::string& prefix) {
  kumo::pollux::functions::sparksql::registerFunctions(prefix);
}

void registerPrestoAggregateFunctions(const std::string& prefix) {
  kumo::pollux::aggregate::prestosql::registerAllAggregateFunctions(prefix);
}

void registerSparkAggregateFunctions(const std::string& prefix) {
  kumo::pollux::functions::aggregate::sparksql::registerAggregateFunctions(
      prefix);
}

exec::AggregateFunctionSignatureMap getAggregateSignatures() {
  return exec::getAggregateFunctionSignatures();
}

void clearAggregateSignatures() {
  exec::aggregateFunctions().withWLock(
      [&](auto& aggregateFunctions) { aggregateFunctions.clear(); });
}

void addSignatureBindings(py::module& m, bool asModuleLocalDefinitions) {
  // TypeSignature
  py::class_<exec::TypeSignature> typeSignature(
      m, "TypeSignature", py::module_local(asModuleLocalDefinitions));
  typeSignature.def("__str__", &exec::TypeSignature::toString);
  typeSignature.def("base_name", &exec::TypeSignature::baseName);
  typeSignature.def("parameters", &exec::TypeSignature::parameters);

  // FunctionSignature
  py::class_<exec::FunctionSignature> functionSignature(
      m, "FunctionSignature", py::module_local(asModuleLocalDefinitions));

  functionSignature.def("__str__", &exec::FunctionSignature::toString);
  functionSignature.def("return_type", &exec::FunctionSignature::returnType);
  functionSignature.def(
      "argument_types", &exec::FunctionSignature::argumentTypes);
  functionSignature.def(
      "variable_arity", &exec::FunctionSignature::variableArity);
  functionSignature.def("variables", &exec::FunctionSignature::variables);
  functionSignature.def(
      "constant_arguments", &exec::FunctionSignature::constantArguments);

  // AggregateFunctionSignature
  py::class_<
      exec::AggregateFunctionSignature,
      std::unique_ptr<exec::AggregateFunctionSignature, py::nodelete>>
      aggregateFunctionSignature(
          m,
          "AggregateFunctionSignature",
          py::module_local(asModuleLocalDefinitions));
  aggregateFunctionSignature.def(
      "__str__", &exec::AggregateFunctionSignature::toString);
  aggregateFunctionSignature.def(
      "intermediate_type", &exec::AggregateFunctionSignature::intermediateType);

  m.def(
      "clear_signatures",
      &clearFunctionRegistry,
      "Clears the function registry.");

  m.def(
      "register_spark_signatures",
      &registerSparkFunctions,
      "Adds Spark signatures to the function registry.",
      py::arg("prefix") = "");

  m.def(
      "register_presto_signatures",
      &registerPrestoFunctions,
      "Adds Presto signatures to the function registry.",
      py::arg("prefix") = "");

  m.def(
      "get_function_signatures",
      []() { return get_function_signatures(); },
      py::return_value_policy::reference,
      "Returns a dictionary of the current signatures.");

  m.def(
      "register_presto_aggregate_signatures",
      &registerPrestoAggregateFunctions,
      "Adds Presto Aggregate signatures to the function registry.",
      py::arg("prefix") = "");

  m.def(
      "register_spark_aggregate_signatures",
      &registerSparkAggregateFunctions,
      "Adds Spark Aggregate signatures to the function registry.",
      py::arg("prefix") = "");

  m.def(
      "get_aggregate_function_signatures",
      &getAggregateSignatures,
      py::return_value_policy::reference,
      "Returns a dictionary of the current aggregate signatures.");

  m.def(
      "clear_aggregate_signatures",
      &clearAggregateSignatures,
      "Clears the Aggregate function registry.");
}
} // namespace kumo::pollux::py
